package com.tianji.service.impl;

import cn.hutool.core.collection.CollUtil;
import cn.hutool.core.util.IdUtil;
import cn.hutool.core.util.StrUtil;
import com.tianji.agent.Agent;
import com.tianji.agent.AgentParams;
import com.tianji.common.utils.UserContext;
import com.tianji.config.ToolResultHolder;
import com.tianji.constants.Constant;
import com.tianji.enums.AgentTypeEnum;
import com.tianji.enums.ChatEventTypeEnum;
import com.tianji.service.ChatService;
import com.tianji.service.ChatSessionService;
import com.tianji.vo.ChatEventVO;
import jakarta.annotation.Resource;
import lombok.extern.slf4j.Slf4j;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.memory.ChatMemory;
import org.springframework.ai.chat.messages.AssistantMessage;
import org.springframework.data.redis.core.StringRedisTemplate;
import reactor.core.publisher.Flux;

import java.util.Map;

@Slf4j
public class AbstractAgent implements Agent {

    @Resource
    private ChatClient chatClient;

    public static final ChatEventVO STOP_EVENT = ChatEventVO.builder().eventType(ChatEventTypeEnum.STOP.getValue()).build();

    @Resource
    private  ChatMemory chatMemory;

    // 存储大模型的生成状态，这里采用ConcurrentHashMap是确保线程安全
    // 目前的版本暂时用Map实现，如果考虑分布式环境的话，可以考虑用redis来实现
    //private static final Map<String, Boolean> GENERATE_STATUS = new ConcurrentHashMap<>();
    @Resource
    private  StringRedisTemplate stringRedisTemplate;

    private static final String GENERATE_STATUS_KEY = "GENERATE_STATUS_KEY";
    @Resource
    private ChatSessionService chatSessionService;

    @Override
    public Flux<ChatEventVO> processStream(String question, String sessionId) {
        // 获取对话id
        String conversationId = ChatService.getConversationId(sessionId);
        var hashOps = stringRedisTemplate.boundHashOps(GENERATE_STATUS_KEY);
        // 大模型输出内容的缓存器，用于在输出中断后的数据存储
        StringBuilder outputBuilder = new StringBuilder();
        // 生成请求id
        String requestId = IdUtil.fastSimpleUUID();
        // 获取当前 用户id
        Long userId = UserContext.getUser();
        AgentParams params = AgentParams.builder()
                .conversationId(conversationId)
                .question(question)
                .userId(userId)
                .requestId(requestId)
                .sessionId(sessionId)
                .build();
        return getChatClientRequestSpec(params)
                .stream()
                .chatResponse()
                .doFirst( ()->{
                    //GENERATE_STATUS.put(sessionId, true);
                    hashOps.put(sessionId, "true");
                })
                .doOnError( throwable ->
                        //GENERATE_STATUS.remove(sessionId)
                        hashOps.delete(sessionId)
                )
                .doOnComplete(() ->
                        //GENERATE_STATUS.remove(sessionId)
                        hashOps.delete(sessionId)
                )
                .doOnCancel(() -> saveStopHistoryRecord(conversationId,outputBuilder.toString()))
                .takeWhile(response -> {
                    return hashOps.get(sessionId) != null;
                })
                .map( chatResponse ->  {

                    // 对于响应结果进行处理，如果是最后一条数据，就把此次消息id放到内存中
                    // 主要用于存储消息数据到 redis中，可以根据消息di获取的请求id，再通过请求id就可以获取到参数列表了
                    // 从而解决，在历史聊天记录中没有外参数的问题
                    var finishReason = chatResponse.getResult().getMetadata().getFinishReason();
                    if (StrUtil.equals(Constant.STOP, finishReason)) {
                        var messageId = chatResponse.getMetadata().getId();
                        ToolResultHolder.put(messageId, Constant.REQUEST_ID, requestId);
                    }
                    // 获取大模型的输出
                    String text = chatResponse.getResult().getOutput().getText();
                    outputBuilder.append(text);
                    return ChatEventVO.builder()
                            .eventData(text)
                            .eventType(ChatEventTypeEnum.DATA.getValue())
                            .build();
                }).concatWith(
                        Flux.defer( ()->{
                            Map<String, Object> map = ToolResultHolder.get(requestId);
                            if (CollUtil.isNotEmpty(map)) {
                                ToolResultHolder.remove(requestId);
                                // 响应给前端参数数据
                                ChatEventVO chatEventVO = ChatEventVO.builder()
                                        .eventType(ChatEventTypeEnum.PARAM.getValue())
                                        .eventData(map)
                                        .build();
                                return Flux.just(chatEventVO, STOP_EVENT);
                            }
                            return Flux.just(STOP_EVENT);
                        })
                ).doFinally(signalType -> {
                    // 异步更新对话信息
                    chatSessionService.update(sessionId,question,userId,outputBuilder.toString());
                });
    }

    private ChatClient.ChatClientRequestSpec getChatClientRequestSpec(AgentParams agentParams) {
        return chatClient.prompt()
                .system(s -> {
                    s.text(this.systemMessage())  // 设置 system prompt
                            .params(this.systemMessageParams(agentParams));  // 设置 system prompt 的参数
                })
                .tools(this.tools()) // 设置 Tools
                .toolContext(this.toolContext(agentParams) // 设置 Tool 上下文参数
                )
                .advisors(advisorSpec -> {
                    advisorSpec.advisors(this.advisors()) // 添加所需的 advisors
                            .params(this.advisorParams(agentParams)); // 给 advisors 设置参数
                })
                .user(agentParams.getQuestion());
    }

    @Override
    public String process(String question, String sessionId) {
        // 获取对话id
        String conversationId = ChatService.getConversationId(sessionId);
        // 生成请求id
        String requestId = IdUtil.fastSimpleUUID();
        // 获取当前 用户id
        Long userId = UserContext.getUser();
        AgentParams params = AgentParams.builder()
                .conversationId(conversationId)
                .question(question)
                .userId(userId)
                .requestId(requestId)
                .sessionId(sessionId)
                .build();
        // 对话
        String content = getChatClientRequestSpec(params)
                .call()
                .content();
        // 更新会话
        chatSessionService.update(sessionId,question,userId,content);
        return content;
    }

    @Override
    public AgentTypeEnum getAgentType() {
        return null;
    }

    private String generateRequestId() {
        return IdUtil.fastSimpleUUID();
    }

    private void saveStopHistoryRecord(String conversationId, String content) {
        chatMemory.add(conversationId, new AssistantMessage(content));
    }


    @Override
    public void stop(String sessionId) {
        stringRedisTemplate.boundHashOps(GENERATE_STATUS_KEY).delete(sessionId);
        //GENERATE_STATUS.remove(sessionId);
    }

    @Override
    public Map<String, Object> advisorParams(AgentParams params) {
        // 将 MessageChatMemoryAdvisor 作为 default advisor 的话 CONVERSATION_ID 参数就是一个通用的
        return Map.of(ChatMemory.CONVERSATION_ID, params.getConversationId());
    }
}